// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <algorithm>
#include <cstdint>

#include <ktest/ktest.h>

#include <nebula/core/memory_pool.h>
#include <tests/core/memory_pool_test.h>
#include <turbo/utility/status.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/version.h>
#include <turbo/log/logging.h>

namespace nebula {

struct DefaultMemoryPoolFactory {
  static MemoryPool* memory_pool() { return default_memory_pool(); }
};

struct SystemMemoryPoolFactory {
  static MemoryPool* memory_pool() { return system_memory_pool(); }
};



template <typename Factory>
class TestMemoryPool : public ::nebula::TestMemoryPoolBase {
 public:
  MemoryPool* memory_pool() override { return Factory::memory_pool(); }
};

TYPED_TEST_SUITE_P(TestMemoryPool);

TYPED_TEST_P(TestMemoryPool, MemoryTracking) { this->TestMemoryTracking(); }

TYPED_TEST_P(TestMemoryPool, OOM) {
#ifndef ADDRESS_SANITIZER
  this->TestOOM();
#endif
}

TYPED_TEST_P(TestMemoryPool, reallocate) { this->TestReallocate(); }

TYPED_TEST_P(TestMemoryPool, Alignment) { this->TestAlignment(); }

REGISTER_TYPED_TEST_SUITE_P(TestMemoryPool, MemoryTracking, OOM, reallocate, Alignment);

INSTANTIATE_TYPED_TEST_SUITE_P(Default, TestMemoryPool, DefaultMemoryPoolFactory);
INSTANTIATE_TYPED_TEST_SUITE_P(System, TestMemoryPool, SystemMemoryPoolFactory);


TEST(DefaultMemoryPool, Identity) {
  // The default memory pool is pointer-identical to one of the backend-specific pools.
  MemoryPool* pool = default_memory_pool();
  std::vector<MemoryPool*> specific_pools = {system_memory_pool()};
  ASSERT_NE(std::find(specific_pools.begin(), specific_pools.end(), pool),
            specific_pools.end());
}

TEST(DefaultMemoryPoolDeathTest, Statistics) {
  MemoryPool* pool = default_memory_pool();
  uint8_t* data1;
  uint8_t* data2;

  ASSERT_OK(pool->allocate(100, &data1));
  ASSERT_OK(pool->allocate(50, &data2));
  pool->free(data2, 50);

  ASSERT_EQ(150, pool->max_memory());
  ASSERT_EQ(150, pool->total_bytes_allocated());
  ASSERT_EQ(100, pool->bytes_allocated());
  ASSERT_EQ(2, pool->num_allocations());

  ASSERT_OK(pool->reallocate(100, 150, &data1));  // Grow data1

  ASSERT_EQ(150, pool->max_memory());
  ASSERT_EQ(150 + 50, pool->total_bytes_allocated());
  ASSERT_EQ(150, pool->bytes_allocated());
  ASSERT_EQ(3, pool->num_allocations());

  ASSERT_OK(pool->reallocate(150, 50, &data1));  // Shrink data1

  ASSERT_EQ(150, pool->max_memory());
  ASSERT_EQ(200, pool->total_bytes_allocated());
  ASSERT_EQ(50, pool->bytes_allocated());
  ASSERT_EQ(3, pool->num_allocations());

  pool->free(data1, 50);

  ASSERT_EQ(150, pool->max_memory());
  ASSERT_EQ(200, pool->total_bytes_allocated());
  ASSERT_EQ(0, pool->bytes_allocated());
  ASSERT_EQ(3, pool->num_allocations());
}

TEST(LoggingMemoryPool, Logging) {
  auto pool = MemoryPool::CreateDefault();

  LoggingMemoryPool lp(pool.get());

  uint8_t* data;
  ASSERT_OK(lp.allocate(100, &data));

  uint8_t* data2;
  ASSERT_OK(lp.allocate(100, &data2));

  lp.free(data, 100);
  lp.free(data2, 100);

  ASSERT_EQ(200, lp.max_memory());
  ASSERT_EQ(200, pool->max_memory());
}

TEST(ProxyMemoryPool, Logging) {
  auto pool = MemoryPool::CreateDefault();

  ProxyMemoryPool pp(pool.get());

  uint8_t* data;
  ASSERT_OK(pool->allocate(100, &data));

  uint8_t* data2;
  ASSERT_OK(pp.allocate(300, &data2));

  ASSERT_EQ(400, pool->bytes_allocated());
  ASSERT_EQ(300, pp.bytes_allocated());

  pool->free(data, 100);
  pp.free(data2, 300);

  ASSERT_EQ(0, pool->bytes_allocated());
  ASSERT_EQ(0, pp.bytes_allocated());
}

}  // namespace nebula
